"""
Description: 创建GelSlim模型并在其表面绘制随机深度场, 仿真三色光照效果
Author: Jin Liu
Date: 2024/06/20
"""
import h5py as h5
import cv2
import tqdm
import numpy as np
from pathlib import Path
from ezgl.items import *
from ezgl import GLViewWidget, tb
from ezgl.items.GLGelSlimItem import GLGelSimItem
from PyQt5 import QtCore, QtGui, QtWidgets
from ezgl.items.sensor import RGBCamera

BASE_DIR = Path(__file__).resolve().parent


class GLView(GLViewWidget):
    def __init__(self, parent=None, **kwargs):
        super().__init__(parent=parent, **kwargs)
        self.setup_ui()

        self.dataset = h5.File(BASE_DIR / 'data/obj_pre.h5', 'a')
        self.length = len(self.dataset["obj_pose"])
        self.idx = -1

        self.left_marker = self.dataset["left_marker"][()]
        self.left_marker = self.left_marker.transpose(0, 2, 3, 1) # TODO

        self.right_marker = self.dataset["right_marker"][()]
        self.right_marker = self.right_marker.transpose(0, 2, 3, 1) # TODO

        self.left_depth = self.dataset["left_depth"]
        self.right_depth = self.dataset["right_depth"]

        # self.dataset.create_dataset("left_rgb", (self.length, 480, 640, 3), maxshape=(None, 480, 640, 3), dtype=np.uint8)
        # self.dataset.create_dataset("right_rgb", (self.length, 480, 640, 3), maxshape=(None, 480, 640, 3), dtype=np.uint8)
        # self.dataset.create_dataset("left_rgb_with_marker", (self.length, 480, 640, 3), maxshape=(None, 480, 640, 3), dtype=np.uint8)
        # self.dataset.create_dataset("right_rgb_with_marker", (self.length, 480, 640, 3), maxshape=(None, 480, 640, 3), dtype=np.uint8)

    def set_img_label(self, idx):
        if idx < 0 or idx >= self.length:
            idx = (idx + self.length) % self.length
        self.idx = idx

        depth = self.right_depth[idx]
        marker = self.right_marker[idx]
        self.model.setDepth(depth*17, marker)

    def sim_img(self, depth, marker, update_shadow):
        if depth is not None:
            depth *= 17
        self.model.setDepth(depth, marker)
        img = self.rgb_camera.render(update_shadow)
        return (img * 255).astype(np.uint8)

    def make_data(self):
        rgb, rgb_with_marker = [], []
        for i in tqdm.trange(self.length):
            depth = self.left_depth[i]
            marker = self.left_marker[i]
            rgb.append(self.sim_img(depth, [], True))
            rgb_with_marker.append(self.sim_img(None, marker, False))
        self.dataset["left_rgb"] = rgb
        self.dataset["left_rgb_with_marker"] = rgb_with_marker

        rgb, rgb_with_marker = [], []
        for i in tqdm.trange(self.length):
            depth = self.right_depth[i]
            marker = self.right_marker[i]
            rgb.append(self.sim_img(depth, [], True))
            rgb_with_marker.append(self.sim_img(None, marker, False))
        self.dataset["right_rgb"] = rgb
        self.dataset["right_rgb_with_marker"] = rgb_with_marker
        print("Done!")

        # for i in tqdm.trange(96, 100):
        #     depth = self.right_depth[i]
        #     marker = self.right_marker[i]
        #     img1 = self.sim_img(depth, [], True)
        #     img2 = self.sim_img(None, marker, False)
        # cv2.imwrite(f"{i}_1.jpg", cv2.cvtColor(img1, cv2.COLOR_RGB2BGR))
        # cv2.imwrite(f"{i}_2.jpg", cv2.cvtColor(img2, cv2.COLOR_RGB2BGR))
        # self.idx = 99
        # print("Done!")

    def setup_ui(self):
        self.camera.set_vector6d((-0.2,  1., -8, -40, -30, -20))

        # 设置光源
        self.light_white = PointLight(pos=[0,0, 1], ambient=(0.2,0.2,0.2), diffuse=(1, 1, 1), directional=True, render_shadow=False, visible=False)
        self.light = PointLight(pos=np.array([-1, -1.73, 1.5])*0.3, diffuse=np.array([0.08, 0.3, 0.7])*1.4*0.8, ambient=(0, 0, 0.),
                                specular=(0, 0, 0.4), linear=0.1, render_shadow=True, visible=False,
                                light_frustum_visible=False)
        self.light1 = PointLight(pos=np.array([2, 0, 1.5])*0.3, diffuse=np.array([0, .7, 0]) * 0.8, ambient=(0, 0., 0),
                                 specular=(0, 0.4, 0), linear=0.1, render_shadow=True, visible=False,
                                 light_frustum_visible=False)
        self.light2 = PointLight(pos=np.array([-1, 1.73, 1.5])*0.3, diffuse=np.array([0.7,0,0])*1.2 * 0.8, ambient=(0., 0, 0),
                                 specular=(0.4, 0, 0), linear=0.1, render_shadow=True, visible=False,
                                 light_frustum_visible=False)

        self.light.setShadow(-0.7, 0.7, -0.7, 0.7, -0.7, 0.7, 0.008, 16)
        self.light1.setShadow(-0.7, 0.7, -0.7, 0.7, -0.7, 0.7, 0.008, 16)
        self.light2.setShadow(-0.7, 0.7, -0.7, 0.7, -0.7, 0.7, 0.008, 16)

        self.scatter = GLScatterPlotItem(size=1, color=(0, 0, 0, 1))

        # 设置坐标轴和地面
        self.ax = GLAxisItem(size=(0.4, 0.4, 0.4), tip_size=0.2).translate(-0.9, -0.9, 0)
        self.grid = GLGridItem(size=(20, 20), spacing=(1, 1), lights=self.light_white).rotate(90, 1, 0, 0).translate(0, 0, -1.5)

        # 设置模型
        self.model = GLGelSimItem(lights=[self.light, self.light1, self.light2]).scale(0.1, 0.1, 0.1)
        self.model.gelslim_gel.translate(0, 0, 0.1)

        gel_material = Material(
                ambient = [0.4, 0.4, 0.4], diffuse = [1.0, 1.0, 1.0],
                specular = [0.2, 0.2, 0.2], shininess = 1., opacity = 1.,
            )
        self.model.gelslim_gel.setMaterial(gel_material)
        self.model.gelslim_base.setMaterial(0, gel_material)

        # 添加 items 到视图
        self.addItems([self.ax, self.model, self.grid])

        # 设置相机传感器
        width, height = 0.674, 0.674*0.75
        self.rgb_camera = RGBCamera(self, eye=(0,0,0.5), up=(-1, 0, 0), ortho_space=(-width,width,-height,height,0,1),
                                    frustum_visible=True)
        # self.depth_camera = DepthCamera(self, eye=(0,0,0.2), up=(-1, 0, 0), ortho_space=(-0.675,0.675,-0.675*0.75,0.675*0.75,-1,1),
        #                                 frustum_visible=False)
        self.rgb_camera.render_group.update(self.model.gelslim_gel, self.model.marker)
        self.light.item_group.update(self.model.gelslim_gel, self.model.marker)
        self.light1.item_group.update(self.model.gelslim_gel, self.model.marker)
        self.light2.item_group.update(self.model.gelslim_gel, self.model.marker)

        # 创建控件
        with tb.window("control", self, 10, pos=(-300, 0), size=(400, 300)):
            default_color = (1.0, 1.0, 1.0, 1.0)
            tb.add_drag_array("RGBA", value=default_color, min_val=0, max_val=1, step=0.01, decimals=2,
                              horizontal=False, format=["R %.2f", "G %.2f", "B %.2f", "A %.2f"],
                              callback=self.on_color_changed)
            tb.add_drag_value("Shininess", value=1, min_val=1, max_val=300, step=1, decimals=0,
                              callback=self.on_color_changed)
            tb.add_separator()
            tb.add_checkbox("render shadow", value=True, callback=self.on_render_shadow)
            with tb.group("iter", show=False):
                tb.add_button("prev", callback=self.on_prev)
                tb.add_button("next", callback=self.on_next)
                tb.add_button("make_data", callback=self.make_data)
            tb.add_timer("onTimeout", 30, callback=self.update, parent=self)
            tb.add_separator()
            with tb.group("Sim Image", collapsible=True, show=True):
                self.sim_image_view = tb.add_image_view("sim", None, img_size=(200, 200))

    def on_next(self):
        self.set_img_label(self.idx+1)
        img = self.rgb_camera.render()
        self.sim_image_view.setData(img*255)
        print(self.idx)

    def on_prev(self):
        self.set_img_label(self.idx-1)
        img = self.rgb_camera.render()
        self.sim_image_view.setData(img*255)
        print(self.idx)

    def on_color_changed(self):
        material = self.model.gelslim_base.getMaterial(0)
        label = self.sender().get_label()
        if label == 'RGBA':
            value = tb.get_value("RGBA")
            material.diffuse = value[:3]
            material.opacity = value[3]
        else:
            material.shininess = tb.get_value("Shininess")

    def on_render_shadow(self, val):
        self.light.set_data(render_shadow=val)
        self.light1.set_data(render_shadow=val)
        self.light2.set_data(render_shadow=val)


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)
    win = GLView(None)
    win.show()
    sys.exit(app.exec_())